import cv2
import numpy as np
import robot


def extract_grid(frame):
    gray = cv2.cvtColor(frame, cv2.COLOR_BGR2GRAY)
    edges = cv2.Canny(gray, 50, 150)
    a, contours, _ = cv2.findContours(edges, cv2.RETR_EXTERNAL, cv2.CHAIN_APPROX_SIMPLE)
    return contours, edges


def is_grid_contour(contour, aspect_ratio_threshold=0.1, min_area=1000):
    x, y, w, h = cv2.boundingRect(contour)
    area = cv2.contourArea(contour)
    aspect_ratio = w / h
    if abs(aspect_ratio - 1) < aspect_ratio_threshold and area > min_area:
        return True
    return False


def find_largest_contour(contours, min_area=5000):
    if contours:
        largest_contour = max(contours, key=cv2.contourArea)
        if cv2.contourArea(largest_contour) >= min_area:
            return largest_contour
    return None


def filter_contours(contours, min_area=5000):
    """
    过滤轮廓，只保留大于指定面积的轮廓。
    :param contours: 轮廓列表
    :param min_area: 最小面积
    :return: 过滤后的轮廓列表
    """
    filtered_contours = [c for c in contours if cv2.contourArea(c) > min_area]
    return filtered_contours


def draw_grid_contour(frame, contour):
    cv2.drawContours(frame, [contour], -1, (0, 0, 255), 2)


def draw_grid_cells(frame, grid_x, grid_y, grid_w, grid_h, rows, cols):
    cell_width = grid_w // cols
    cell_height = grid_h // rows
    for row in range(rows):
        for col in range(cols):
            start_x = grid_x + col * cell_width
            start_y = grid_y + row * cell_height
            end_x = start_x + cell_width
            end_y = start_y + cell_height
            cv2.rectangle(frame, (start_x, start_y), (end_x, end_y), (0, 255, 255), 1)


def detect_black_mask(frame):
    hsv = cv2.cvtColor(frame, cv2.COLOR_BGR2HSV)
    lower_black = np.array([0, 0, 200])
    upper_black = np.array([180, 50, 255])
    mask = cv2.inRange(hsv, lower_black, upper_black)
    kernel = np.ones((5, 5), np.uint8)
    mask = cv2.morphologyEx(mask, cv2.MORPH_CLOSE, kernel)
    mask = cv2.morphologyEx(mask, cv2.MORPH_OPEN, kernel)
    cv2.imshow('white Mask', mask)
    return mask


def get_grid_position(x, y, grid_x, grid_y, grid_w, grid_h, rows, cols):
    cell_width = grid_w // cols
    cell_height = grid_h // rows
    col = (x - grid_x) // cell_width
    row = (y - grid_y) // cell_height
    row = min(max(row, 0), rows - 1)
    col = min(max(col, 0), cols - 1)
    return row, col


def smooth_contours(contours, alpha=0.5):
    smoothed_contours = []
    for contour in contours:
        # Apply smoothing (e.g., Gaussian Blur)
        contour_smoothed = cv2.GaussianBlur(contour, (5, 5), 0)
        smoothed_contours.append(contour_smoothed)
    return smoothed_contours


first_frame = None
second_frame = None
first_mask = None
second_mask = None
white_row = None
white_col = None
last_row = None
last_col = None
robot_start_x = None
robot_start_y = None


def main_function(frame, key,start_row,start_col):
    global first_frame
    global second_frame
    global first_mask
    global second_mask
    global white_row
    global white_col
    global last_row
    global last_col

    robot_put_inedx = None
    contours, edges = extract_grid(frame)

    if key == ord('1'):
        first_frame = frame.copy()
        first_mask = detect_black_mask(first_frame)
        print("First frame recorded.")

    elif key == ord('2'):
        second_frame = frame.copy()
        second_mask = detect_black_mask(second_frame)
        print("Second frame recorded.")

    if first_frame is not None and second_frame is not None:
        largest_contour = find_largest_contour(contours, min_area=15000)
        if largest_contour is not None:
            x, y, w, h = cv2.boundingRect(largest_contour)
            draw_grid_contour(frame, largest_contour)
            draw_grid_cells(frame, x, y, w, h, rows=3, cols=3)

            cell_width = w // 3
            cell_height = h // 3

            # 创建一个mask以识别在第二帧中出现的新区域
            new_area_mask = np.zeros_like(second_mask)

            # 在第一帧掩膜中找出新区域
            for row in range(3):
                for col in range(3):
                    start_x = x + col * cell_width
                    start_y = y + row * cell_height
                    end_x = start_x + cell_width
                    end_y = start_y + cell_height

                    # 提取当前单元格的掩膜
                    cell_mask_first = first_mask[start_y:end_y, start_x:end_x]
                    cell_mask_second = second_mask[start_y:end_y, start_x:end_x]

                    # 计算新区域的掩膜
                    new_mask = cv2.bitwise_and(cell_mask_second, cv2.bitwise_not(cell_mask_first))
                    new_area_mask[start_y:end_y, start_x:end_x] = new_mask

                    # 计算新区域的面积
                    new_area = np.sum(new_mask == 255)

                    if new_area > 1000:
                        # print(f"New area in Cell at Row: {row}, Col: {col} is larger than 1000 pixels.")

                        # 绘制单元格
                        cv2.rectangle(frame, (start_x, start_y), (end_x, end_y), (0, 255, 0), 2)
                        cv2.putText(frame, f'Row: {row}, Col: {col}', (start_x, start_y - 10),
                                    cv2.FONT_HERSHEY_SIMPLEX, 0.5, (0, 255, 0), 1)

                        if row is not None and col is not None:
                            if row != last_row or col != last_col:
                                last_row = row
                                last_col = col

                                output_row, output_col = robot.main_function(row, col,start_row,start_col)
                                robot_put_inedx = output_row * 3 + output_col

    if edges is not None:
        # cv2.imshow('Edge Detection', edges)
        pass

    return frame,robot_put_inedx
